import cv2
import numpy as np
from ultralytics import YOLO
import os
import glob
import warnings



class trunk_road_detect():
    def __init__(self):
        pass
        #self.name_list = [] 
    def pridict(self,img_path,show_yolov8Label = True):
        if self.name_list == []:
            print("First load the model by running load_wights")
            return None
        origin_img = cv2.imread(img_path)
        result = self.detect(origin_img,show=False,verbose=False)
        if show_yolov8Label:
            img = result.plot()
        else:
            img =   origin_img
        trunk_left,trunk_right = self.get_boxes(result,origin_img)
        if len(trunk_left)<2 or len(trunk_right)<2:
            return
        self.draw_road(img,trunk_left,trunk_right,img_path)

    def load_wights(self,path_wights):
        self.model = YOLO(path_wights)
        self.name_list = self.model.names
    def detect(self,path_img,show=False,verbose=False):
        results = self.model.predict(path_img,show=show,verbose=verbose)
        result = results[0]
        return result
    def get_boxes(self,result,origin_img):
        if self.name_list ==None:
            return None
        name_list = self.name_list
        sort_list = result.boxes.cls.tolist()  
        trunk_left = []
        trunk_right = []
        for i in range(len(sort_list)):
            if name_list[sort_list[i]] == "trunk":  
                trunk_info = result.boxes[i].xywh.cpu().numpy()[0]
                x, y, w, h = trunk_info
                if x < origin_img.shape[1] / 2:
                    trunk_left.append([x, y + h / 2])
                else:
                    trunk_right.append([x, y + h / 2])
        return [trunk_left,trunk_right]
    def draw_road(self,img,trunk_left,trunk_right,img_path=None):
        trunk_left = np.array(sorted(trunk_left, key=lambda x: x[0]))  
        trunk_right = np.array(sorted(trunk_right, key=lambda x: x[0]))
        warnings.filterwarnings("ignore", category=np.RankWarning)
        line_left = np.polyfit(trunk_left[:, 0], trunk_left[:, 1], 1)  
        line_right = np.polyfit(trunk_right[:, 0], trunk_right[:, 1], 1)  
       
        warnings.filterwarnings("default")
   
        height, width, _ = img.shape
        half_width = int(width / 2)
     
        point_left = {"x1":0,
                      "y1":int(line_left[1]),
                      "x2":half_width - 1,
                      "y2":int(line_left[0] * (half_width - 1) + line_left[1]),
                      }
        point_right = {"x1":half_width,
                       "y1":int(line_right[0] * (half_width) + line_right[1]),
                       "x2":width - 1,
                       "y2":int(line_right[0] * (width - 1) + line_right[1]),
                       }
     
        if point_left["y2"] > point_right["y1"]:
            point_right["y1"] = point_left["y2"]
            point_right["x1"] =int((point_right["y1"] - line_right[1])/line_right[0])
        else:
            point_left["y2"] = point_right["y1"]
            point_left["x2"] = int((point_left["y2"] - line_left[1]) / line_left[0])
   
        cv2.line(img, (point_left["x1"], point_left["y1"]), (point_left["x2"], point_left["y2"]), (0, 0, 255), 2)
        cv2.line(img, (point_right["x1"], point_right["y1"]), (point_right["x2"], point_right["y2"]), (0, 0, 255), 2)
       
        cv2.line(img, (int((point_left["x1"]+point_right["x2"])/2), int((point_left["y1"]+point_right["y2"])/2)),
                 (int((point_left["x2"]+point_right["x1"])/2), int((point_left["y2"]+point_right["y1"])/2)),
                 (0, 255, 255),2)
        cv2.imshow("img", img)
        save_path = os.path.join("out",img_path)
        cv2.imwrite(save_path,img)
        cv2.waitKey(1)    
        return img


if __name__ == "__main__":
    detect_app = trunk_road_detect()
    detect_app.load_wights(r"weights/")
    img_list = glob.glob(r"images/*.jpg")
    #detect_app.pridict(img_list[0])
    for _img in img_list:
        detect_app.pridict(_img)
